package net.xuele.learn.flink.book;

import net.xuele.learn.flink.book.utils.SensorReading;
import net.xuele.learn.flink.book.utils.SensorSource;
import net.xuele.learn.flink.book.utils.ThresholdUpdate;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction;
import org.apache.flink.util.Collector;

/**
 * @Author patrick
 * @Date 2023/7/12 9:21
 * @Description： 下面的代码是KeyedBroadcastProcessFunction函数的应用，它在processBroadcastElement方法中将广播的状态进行了按照key的分发
 *
 * 主要做法在于广播数据和源数据之间存在相同的key值，在此例中，即id，因此实现了将广播值按照key分发
 */
public class KeyedBroadcast {

    public static void main(String[] args) {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStream<SensorReading> sensorData = env.addSource(new SensorSource());

        ThresholdUpdate thresholdUpdate = ThresholdUpdate.builder().id("sensor_1").threshold(10.0).build();
        DataStream<ThresholdUpdate> thresholds = env.fromElements(thresholdUpdate);

        KeyedStream<SensorReading, String> keyedSensorData = sensorData.keyBy(value -> value.id);

        MapStateDescriptor<String, Double> mapStateDescriptor = new MapStateDescriptor<>("thresholds", String.class, Double.class);

        BroadcastStream<ThresholdUpdate> broadcast = thresholds.broadcast(mapStateDescriptor);

        keyedSensorData.connect(broadcast).process(new UpdatableTemperatureAlertFunction())
                .print();
    }

    private static class UpdatableTemperatureAlertFunction extends KeyedBroadcastProcessFunction<String, SensorReading,
            ThresholdUpdate, Tuple3<String, Double, Double>> {

        private static final long serialVersionUID = -4336601595944399761L;

        private final MapStateDescriptor<String, Double> mapStateDescriptor = new MapStateDescriptor<>("thresholds", String.class, Double.class);

        private ValueState<Double> lastTempState;

        @Override
        public void open(Configuration parameters) throws Exception {
            ValueStateDescriptor<Double> lastTempDescriptor = new ValueStateDescriptor<>("lastTempState", Double.class);
            lastTempState = getRuntimeContext().getState(lastTempDescriptor);

        }

        /**
         * 由于不存在跨任务的通信，所以只有processBroadcastElement方法才可以更新广播状态。
         *      用于确保所有井行实例中广播状态所保存的信息完全相同
         * @param value
         * @param ctx  ReadOnlyContext，不能更新广播状态
         * @param out
         * @throws Exception
         */
        @Override
        public void processElement(SensorReading value, ReadOnlyContext ctx, Collector<Tuple3<String, Double, Double>> out) throws Exception {
            ReadOnlyBroadcastState<String, Double> thresholds = ctx.getBroadcastState(mapStateDescriptor);
            // 根据二者相同的key检查阈值是否存在
            if (thresholds.contains(value.id)) {
                Double threshold = thresholds.get(value.id);
                Double lastTemp = lastTempState.value();
                if (lastTemp != null) {
                    double tempDiff = Math.abs(value.temperature - lastTemp);
                    // 温度差值超过阈值，需要告警
                    if (tempDiff > threshold) {
                        out.collect(Tuple3.of(value.id, value.temperature, tempDiff));
                    }
                }
            }
            lastTempState.update(value.temperature);
        }

        /**
         * applyToKeyedState() allows to register a KeyedStateFunction to be applied to all states of all keys
         * associated with the provided stateDescriptor
         * @param value
         * @param ctx
         * @param out
         * @throws Exception
         */
        @Override
        public void processBroadcastElement(ThresholdUpdate value, Context ctx, Collector<Tuple3<String, Double, Double>> out) throws Exception {
            BroadcastState<String, Double> thresholds = ctx.getBroadcastState(mapStateDescriptor);
            if (value.threshold != 0.0) {
                thresholds.put(value.id, value.threshold);
            } else {
                thresholds.remove(value.id);
            }
        }
    }
}
